import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report, accuracy_score
import torch
from torch import nn
import torch.utils.data as Data
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
from torchvision.datasets import CIFAR10
from torchvision.utils import make_grid
from torchinfo import summary
from tqdm import tqdm
from ae_net import AE_FC, AE_CNN, VAE_FC, VAE_CNN
import torchvision
import torchvision.datasets as dset
from train import train_model, train_vae_model
from util import get_project_root

Root_path = get_project_root()

def train_AE_CNN():
    #生成模型
    image_size=64
    # net = AE_FC(image_size, 64)
    net = AE_CNN()
    summary(net, (1, 3, image_size, image_size))

    # trans = transforms.Compose((transforms.Resize(32), transforms.ToTensor()))
    trans = transforms.Compose((transforms.RandomCrop(image_size), transforms.ToTensor()))
    back_size = 1
    batch_size = 900

    dataset_name = "pathlo"
    trainset = dset.ImageFolder(root=Root_path + "/data/pathol256", transform=trans)

    # optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)  # 选好优化方式
    optimizer = torch.optim.RMSprop(net.parameters(), lr=1e-5)
    # 定义损失函数和优化方式
    loss_fn = nn.MSELoss(reduction='sum')  # 损失函数为交叉熵，多用于多分类问题 reduction='sum'

    train_model(net, dataset_name, trainset, optimizer,loss_fn, back_size, batch_size)

def train_VAE_CNN():
    #生成模型
    image_size=64
    net = VAE_CNN(image_size,)
    summary(net, (1, 3, image_size, image_size))

    trans = transforms.Compose((transforms.RandomCrop(image_size), transforms.ToTensor()))
    back_size = 1
    batch_size = 256

    dataset_name = "pathlo"
    trainset = dset.ImageFolder(root=Root_path + "/data/pathol256", transform=trans)

    # optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)  # 选好优化方式
    # optimizer = torch.optim.SGD(net.parameters(), lr=1e-5)  # 选好优化方式
    optimizer = torch.optim.RMSprop(net.parameters(), lr=1e-5)  # 选好优化方式

    loss_fn = nn.MSELoss(reduction='sum')  # 损失函数为交叉熵，多用于多分类问题 reduction='sum' mean

    train_vae_model(net, dataset_name, trainset, optimizer,loss_fn, back_size, batch_size)

if __name__ == '__main__':
    train_VAE_CNN()